# 安装第三方库： pip install pymysql hashlib
import pymysql
import hashlib


# 建立并返回一个MySQL数据库连接
def get_connection():
    return pymysql.connect(
        host="localhost",
        port=3306,
        user="root",
        password="123456",
        database="demo_db",
        charset="utf8mb4",
    )


# 建用户表users（如不存在），包含id、username、password、realname、is_frozen字段。
def create_user_table():
    conn = get_connection()
    try:
        with conn.cursor() as cursor:
            cursor.execute(
                """
                CREATE TABLE IF NOT EXISTS users (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    username VARCHAR(50) UNIQUE NOT NULL,
                    password VARCHAR(64) NOT NULL,
                    realname VARCHAR(50) NOT NULL,
                    is_frozen BOOLEAN DEFAULT 0
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
            """
            )
        conn.commit()
    finally:
        conn.close()


# 对输入的明文密码进行SHA-256加密，返回加密后的十六进制字符串。
def hash_password(password):
    return hashlib.sha256(password.encode("utf-8")).hexdigest()


# 向users表插入初始用户数据（如已存在则跳过）。
def insert_initial_users():
    users = [
        ("user1", hash_password("123456"), "张三", 0),
        ("user2", hash_password("123456"), "李四", 0),
        ("user3", hash_password("123456"), "王五", 1),
    ]
    conn = get_connection()
    try:
        with conn.cursor() as cursor:
            for u in users:
                try:
                    cursor.execute(
                        "INSERT INTO users (username, password, realname, is_frozen) VALUES (%s, %s, %s, %s)",
                        u,
                    )
                except pymysql.err.IntegrityError:
                    pass  # 已存在则跳过
        conn.commit()
    finally:
        conn.close()


# 注册新用户，用户名唯一，密码加密存储，注册成功或用户名已存在时给出提示。
def register(username, password, realname):
    conn = get_connection()
    try:
        with conn.cursor() as cursor:
            try:
                cursor.execute(
                    "INSERT INTO users (username, password, realname) VALUES (%s, %s, %s)",
                    (username, hash_password(password), realname),
                )
                conn.commit()
                print("注册成功！")
            except pymysql.err.IntegrityError:
                print("用户名已存在！")
    finally:
        conn.close()


# 用户登录，校验账号是否存在、是否被冻结、密码是否正确，并给出相应提示。
def login(username, password):
    conn = get_connection()
    try:
        with conn.cursor() as cursor:
            cursor.execute(
                "SELECT password, realname, is_frozen FROM users WHERE username=%s",
                (username,),
            )
            row = cursor.fetchone()
            if not row:
                print("账号不存在！")
                return
            db_password, realname, is_frozen = row
            if is_frozen:
                print("账号已被冻结！")
                return
            if db_password == hash_password(password):
                print(f"登录成功，欢迎 {realname}！")
            else:
                print("密码错误！")
    finally:
        conn.close()


# 程序主入口，提供注册和登录操作菜单。
def main():
    create_user_table()
    insert_initial_users()
    print("1. 注册\n2. 登录")
    choice = input("请选择操作(1/2): ")
    if choice == "1":
        username = input("请输入账号: ")
        password = input("请输入密码: ")
        realname = input("请输入姓名: ")
        register(username, password, realname)
    elif choice == "2":
        username = input("请输入账号: ")
        password = input("请输入密码: ")
        login(username, password)
    else:
        print("无效选择！")


if __name__ == "__main__":
    main()
